import json
from typing import Iterable

from subway_qa.adapters import chat

PROMPT = """
你是地铁线路查询智能客服

你的任务是从与用户的对话记录中抽取出用户要查询的地铁线路的起点和终点，并通过json格式进行返回。如果对话记录与地铁线路查询无关，则无需抽取。

对话记录1（用户直接说明起点终点）：
User: 你好我想查地铁线路
Assistant: 好的，请问您要查询从哪到哪的地铁线路？
User: 从西湖文化广场到火车东站
请抽取对话记录中的地铁线路查询信息，并返回json格式。

{ 
  "start": "西湖文化广场",
  "end": "火车东站",
  "is_relevant": true
}

对话记录2（用户只说了起点）:
User: 从武林广场出发怎么坐地铁
请抽取对话记录中的地铁线路查询信息，并返回json格式。

{ 
  "start": "武林广场",
  "end": "",
  "is_relevant": true
}

对话记录3（用户提出无关请求）：
User: 你好
Assistant: 你好，请问有什么可以帮您的吗
User: 噢噢不好意思，你们这里可以办公交卡么
请抽取对话记录中的地铁线路查询信息，并返回json格式。

{ 
  "start": "",
  "end": "",
  "is_relevant": false
}

对话记录4（用户说了换乘需求）：
User: 我想查地铁线路
Assistant: 好的，请问您要查询从哪到哪的地铁线路？
User: 从城站到钱江世纪城，要换乘的
请抽取对话记录中的地铁线路查询信息，并返回json格式。

{ 
  "start": "城站",
  "end": "钱江世纪城",
  "is_relevant": true
}

对话记录5（用户前面是地铁查询，但最后是无关问题）：
User: 你好我想查地铁
Assistant: 好的，请问您要查询从哪到哪的地铁线路？
User: 从凤起路到龙翔桥
Assistant: 为您查询到以下线路...
User: 谢谢，那公交卡在哪里充值？
请抽取对话记录中的地铁线路查询信息，并返回json格式。


{ 
  "start": "凤起路",
  "end": "龙翔桥",
  "is_relevant": false
}

对话记录6（用户说了查找地铁线路，但最后是无关问题）：
User: 我想查地铁线路
请抽取对话记录中的地铁线路查询信息，并返回json格式。

{ 
  "start": "",
  "end": "",
  "is_relevant": true
}

对话记录7（问线路规划）：
User：知春路到枣园怎么走
请抽取对话记录中的地铁线路查询信息，并返回json格式。

{ 
  "start": "知春路",
  "end": "枣园",
  "is_relevant": true
}

上述是一些抽取示例，仅供参考
"""


def extract(messages: list[dict]) -> dict:
    text = ""
    for message in messages:
        if not isinstance(message["content"], Iterable):
            continue
        if message["role"] == "user":
            text += f"User: {message['content'][0]}\n"
        else:
            text += f"Assistant: {message['content'][0]}\n"

    text += "请抽取对话记录中的地铁线路查询信息，并返回json格式。"
    messages = [
        {"role": "system", "content": [{"type": "text", "text": PROMPT}]},
        {"role": "user", "content": [{"type": "text", "text": text}]},
    ]
    response_text = chat(messages)

    try:
        start = response_text.find("{")
        end = response_text.rfind("}") + 1
        return json.loads(response_text[start:end])
    except Exception as e:
        pass

    try:
        start = response_text.find("[")
        end = response_text.rfind("]") + 1
        return json.loads(response_text[start:end])[-1]
    except Exception as e:
        pass
    return {"start": "", "end": "", "is_relevant": False}


if __name__ == "__main__":
    from subway_qa.adapters import chat

    print(
        extract(
            [
                {"role": "user", "content": "你好我想查地铁线路"},
            ]
        )
    )